import sys, os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import json
import re
import numpy as np
import pandas as pd

from preprossData.utils.crop_padding import crop_resize_save

# 诊断映射规则
diagnosis_mapping = {
    0: "normal",
    1: "glaucoma",
    2: "suspicious glaucoma"
}

# 性别映射规则：0 -> Male, 1 -> Female
gender_mapping = {
    0: "Male",
    1: "Female"
}

def clean_field(field):
    """
    清理字段名称，去除括号及括号内内容，并转为小写
    """
    return re.sub(r'\s*\([^)]*\)', '', field).strip().lower()

def parse_info_xlsx(xlsx_path):
    """
    读取 xlsx 文件，并返回一个字典，结构为：
    
    {
      "<patient_id>": {
            "patient info": { ... 其他临床信息 ... },
            "text": "<diagnosis text>"
      },
      ...
    }
    
    说明：
      - Gender 字段按 0->Male, 1->Female 转换
      - Diagnosis 数值根据映射转换为文本：0->normal, 1->glaucoma, 2->suspicious glaucoma
      - 字段名称统一转为小写，并去除括号内的内容
      - 排除的字段：device, plus form, diagnosis code, id, serie number, diagnosis
      - 若某个字段值为 None 或 NaN，则不存储该键
    """
    result = {}
    df = pd.read_excel(xlsx_path, engine='openpyxl')
    df.columns = [clean_field(col) for col in df.columns]
    
    exclude = {"device", "plus form", "diagnosis code", "id", "serie number", "diagnosis"}
    
    for _, row in df.iterrows():
        # 获取患者 ID，假设字段名为 "id"
        pid = str(row["id"]).strip()
        if pid[0]=='#':
            pid=pid[1:] # 先去除#统一处理
        # 构建诊断文本，仅使用 diagnosis 字段
        dg_val = row.get("diagnosis", None)
        
        diag_key = int(dg_val)
        diag_text = diagnosis_mapping.get(diag_key)
        
        info = {}
        for key, value in row.to_dict().items():
            if key in exclude:
                continue
            if value is None or (isinstance(value, float) and np.isnan(value)):
                continue
            # 特殊处理 gender 字段
            if key == "gender":
                try:
                    g_val = int(value)
                except (ValueError, TypeError):
                    g_val = value
                info[key] = gender_mapping.get(g_val, value)
            else:
                info[key] = value
        result[pid] = {"patient info": info, "text": diag_text}
    return result

def merge_patient_data(dir_path):
    """
    读取指定目录下的 patient_data_os.xlsx 和 patient_data_od.xlsx，
    为 OS 数据添加 "OS" 后缀，为 OD 数据添加 "OD" 后缀，
    并在 key 前加上 '#'，合并为一个完整的 dict 返回。
    
    例如：
    "#018OS": {'patient info': { ... }, 'text': 'suspicious glaucoma'}
    """
    merged = {}
    file_os = os.path.join(dir_path, "patient_data_os.xlsx")
    file_od = os.path.join(dir_path, "patient_data_od.xlsx")
    
    data_os = parse_info_xlsx(file_os)
    for pid, info in data_os.items():
        new_key = f"#{pid}OS"
        merged[new_key] = info
        
    data_od = parse_info_xlsx(file_od)
    for pid, info in data_od.items():
        new_key = f"#{pid}OD"
        merged[new_key] = info
        
    return merged

def gather_data(data_path, tar_path, prefix='papila_'):
    """
    对数据集进行预处理：
      - 数据集根目录中要求包含 ClinicalData 文件夹（存放 patient_data_os.xlsx 和 patient_data_od.xlsx）
        以及 FundusImages 文件夹（存放所有原始图像）。
      - 先通过 merge_patient_data 解析 ClinicalData 文件夹中的临床信息，
        生成一个 info_dict，其 key 格式为 "#<patient_id>OS" 或 "#<patient_id>OD"。
      - 遍历 FundusImages 文件夹中的所有图像，利用文件名的特定位置提取 key，
        其中 key = f"#{image_name[3:-4]}"（请勿修改该提取方式）。
      - 对每张图像调用 crop_resize_save 进行裁剪与 resize，保存为 PNG 格式，
        新图像名称为 {prefix}_{原始文件名}.png（保留原始文件名，仅修改后缀并加上前缀）。
      - 将每张图像对应的标注信息写入 data_dict，包括 patient_info 和 diagnosis 的 text 字段，
        最终将 data_dict 保存到 tar_path/annotations.json 中。
    
    参数：
        data_path (str): 数据集根目录，要求其中包含 ClinicalData 和 FundusImages 文件夹。
        tar_path (str): 预处理后数据存放目录。
        prefix (str): 处理后图片名称的前缀，默认为 "papila_".
    
    返回：
        dict: 标注信息字典。
    """
    os.makedirs(tar_path, exist_ok=True)
    images_out_dir = os.path.join(tar_path, 'images')
    os.makedirs(images_out_dir, exist_ok=True)
    
    # 解析临床数据
    info_dict = merge_patient_data(os.path.join(data_path, "ClinicalData"))
    
    image_src_dir = os.path.join(data_path, 'FundusImages')
    data_dict = {}
    
    # 遍历 FundusImages 下的所有图像（假设为平铺存放）
    for image_name in sorted(os.listdir(image_src_dir)):
        # 提取 key，注意不要修改该提取方式
        k = f"#{image_name[3:-4]}"
        if k not in info_dict:
            print(f"Warning: key {k} not found in clinical info. Skipping {image_name}.")
            continue
        # 从 info_dict 中获取患者信息与诊断文本
        patient_info = info_dict[k].get("patient info", {})
        diagnosis_text = info_dict[k].get("text", "")
        
        # 保留原始文件名，仅修改后缀并加上前缀
        base_name = os.path.splitext(image_name)[0]
        new_image_name = f"{prefix}_{base_name}.png"
        dest_image_path = os.path.join(images_out_dir, new_image_name)
        
        # 对图像进行裁剪与 resize，保存为 PNG 格式
        crop_info=crop_resize_save(
            image_path=os.path.join(image_src_dir, image_name),
            save_path=dest_image_path,
            resize=(224, 224),
            crop_threshold=25
        )
        
        data_dict[new_image_name] = {
            "image_name": new_image_name,
            "image_path": os.path.join("images", new_image_name),
            'original_path': os.path.relpath(os.path.join(image_src_dir, image_name), data_path),
            'crop_info':crop_info,   
            "patient_info": patient_info,
            "diagnosis": {
                "text": diagnosis_text
            }
        }
    
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(data_dict, f, indent=4)
    
    return data_dict

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description="Infant Retinal Database 新数据集预处理")
    parser.add_argument("--data_path", type=str, default="/mnt/e/Document_Workspace/Dataset/public_dataset/PAPILA/",
                        help="新数据集根目录，要求其下包含类别文件夹")
    parser.add_argument("--tar_path", type=str, default="/mnt/e/Document_Workspace/Dataset/processed224_new/PAPILA",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="papila_",
                        help="处理后图片名称的前缀，默认为 'papila_'")
    args = parser.parse_args()
    
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix)
    print("WARNING: 本数据集的标注是按照一次诊断的标注，因此，虽然部分图片被标注为了阳性，但不一定能在该图片观察到，可能在同一次随访拍摄的其他图片中观察到病灶。")
    print("Preprocessing completed. Annotations saved.")
